// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//

// This API is EXPERIMENTAL.

#pragma once

#include <memory>
#include <string>
#include <utility>
#include <variant>
#include <vector>

#include <nebula/compute/type_fwd.h>
#include <nebula/core/datum.h>
#include <nebula/types/type_fwd.h>
#include <turbo/container/small_vector.h>

namespace nebula::compute {

    /// \defgroup expression-core Expressions to describe data transformations
    ///
    /// @{

    /// An unbound expression which maps a single Datum to another Datum.
    /// An expression is one of
    /// - A literal Datum.
    /// - A reference to a single (potentially nested) field of the input Datum.
    /// - A call to a compute function, with arguments specified by other Expressions.
    class TURBO_EXPORT Expression {
    public:
        struct Call {
            std::string function_name;
            std::vector<Expression> arguments;
            std::shared_ptr<FunctionOptions> options;
            // Cached hash value
            size_t hash;

            // post-Bind properties:
            std::shared_ptr<Function> function;
            const Kernel *kernel = nullptr;
            std::shared_ptr<KernelState> kernel_state;
            TypeHolder type;

            void ComputeHash();
        };

        [[nodiscard]] std::string to_string() const;

        [[nodiscard]] bool equals(const Expression &other) const;

        [[nodiscard]] size_t hash() const;

        struct Hash {
            size_t operator()(const Expression &expr) const { return expr.hash(); }
        };

        /// Bind this expression to the given input type, looking up Kernels and field types.
        /// Some expression simplification may be performed and implicit casts will be inserted.
        /// Any state necessary for execution will be initialized and returned.
        turbo::Result<Expression> Bind(const TypeHolder &in, ExecContext * = nullptr) const;

        turbo::Result<Expression> Bind(const Schema &in_schema, ExecContext * = nullptr) const;

        // XXX someday
        // Clone all KernelState in this bound expression. If any function referenced by this
        // expression has mutable KernelState, it is not safe to execute or apply simplification
        // passes to it (or copies of it!) from multiple threads. Cloning state produces new
        // KernelStates where necessary to ensure that Expressions may be manipulated safely
        // on multiple threads.
        // turbo::Result<ExpressionState> CloneState() const;
        // turbo::Status set_state(ExpressionState);

        /// Return true if all an expression's field references have explicit types
        /// and all of its functions' kernels are looked up.
        [[nodiscard]] bool is_bound() const;

        /// Return true if this expression is composed only of Scalar literals, field
        /// references, and calls to ScalarFunctions.
        [[nodiscard]] bool is_scalar_expression() const;

        /// Return true if this expression is literal and entirely null.
        [[nodiscard]] bool is_null_literal() const;

        /// Return true if this expression could evaluate to true. Will return true for any
        /// unbound or non-boolean Expressions. is_satisfiable does not (currently) do any
        /// canonicalization or simplification of the expression, so even Expressions
        /// which are unsatisfiable may spuriously return `true` here. This function is
        /// intended for use in predicate pushdown where a filter expression is simplified
        /// by a guarantee, so it assumes that trying to simplify again would be redundant.
        [[nodiscard]] bool is_satisfiable() const;

        // XXX someday
        // turbo::Result<PipelineGraph> GetPipelines();

        [[nodiscard]] bool is_valid() const { return impl_ != nullptr; }

        /// Access a Call or return nullptr if this expression is not a call
        [[nodiscard]] const Call *call() const;

        /// Access a Datum or return nullptr if this expression is not a literal
        [[nodiscard]] const Datum *literal() const;

        /// Access a FieldRef or return nullptr if this expression is not a field_ref
        [[nodiscard]] const FieldRef *field_ref() const;

        /// The type to which this expression will evaluate
        [[nodiscard]] const DataType *type() const;
        // XXX someday
        // NullGeneralization::type nullable() const;

        struct Parameter {
            FieldRef ref;

            // post-bind properties
            TypeHolder type;
            turbo::SmallVector<int, 2> indices;
        };

        [[nodiscard]] const Parameter *parameter() const;

        Expression() = default;

        explicit Expression(Call call);

        explicit Expression(Datum literal);

        explicit Expression(Parameter parameter);

    private:
        using Impl = std::variant<Datum, Parameter, Call>;
        std::shared_ptr<Impl> impl_;

        friend bool Identical(const Expression &l, const Expression &r);
    };

    inline bool operator==(const Expression &l, const Expression &r) { return l.equals(r); }

    inline bool operator!=(const Expression &l, const Expression &r) { return !l.equals(r); }

    TURBO_EXPORT void print_to(const Expression &, std::ostream *);

    // Factories

    TURBO_EXPORT
    Expression literal(Datum lit);

    template<typename Arg>
    Expression literal(Arg &&arg) {
        return literal(Datum(std::forward<Arg>(arg)));
    }

    TURBO_EXPORT
    Expression field_ref(FieldRef ref);

    TURBO_EXPORT
    Expression call(std::string function, std::vector<Expression> arguments,
                    std::shared_ptr<FunctionOptions> options = nullptr);

    template<typename Options, typename = typename std::enable_if<
            std::is_base_of<FunctionOptions, Options>::value>::type>
    Expression call(std::string function, std::vector<Expression> arguments,
                    Options options) {
        return call(std::move(function), std::move(arguments),
                    std::make_shared<Options>(std::move(options)));
    }

    /// Assemble a list of all fields referenced by an Expression at any depth.
    TURBO_EXPORT
    std::vector<FieldRef> fields_in_expression(const Expression &);

    /// Check if the expression references any fields.
    TURBO_EXPORT
    bool expression_has_field_refs(const Expression &);

    struct TURBO_EXPORT KnownFieldValues;

    /// Assemble a mapping from field references to known values. This derives known values
    /// from "equal" and "is_null" Expressions referencing a field and a literal.
    TURBO_EXPORT
    turbo::Result<KnownFieldValues> extract_known_field_values(
            const Expression &guaranteed_true_predicate);

    /// @}

    /// \defgroup expression-passes Functions for modification of Expressions
    ///
    /// @{
    ///
    /// These transform bound expressions. Some transforms utilize a guarantee, which is
    /// provided as an Expression which is guaranteed to evaluate to true. The
    /// guaranteed_true_predicate need not be bound, but canonicalization is currently
    /// deferred to producers of guarantees. For example in order to be recognized as a
    /// guarantee on a field value, an Expression must be a call to "equal" with field_ref LHS
    /// and literal RHS. Flipping the arguments, "is_in" with a one-long value_set, ... or
    /// other semantically identical Expressions will not be recognized.

    /// Weak canonicalization which establishes guarantees for subsequent passes. Even
    /// equivalent Expressions may result in different canonicalized expressions.
    /// TODO this could be a strong canonicalization
    TURBO_EXPORT
    turbo::Result<Expression> canonicalize(Expression, ExecContext * = nullptr);

    /// Simplify Expressions based on literal arguments (for example, add(null, x) will always
    /// be null so replace the call with a null literal). Includes early evaluation of all
    /// calls whose arguments are entirely literal.
    TURBO_EXPORT
    turbo::Result<Expression> fold_constants(Expression);

    /// Simplify Expressions by replacing with known values of the fields which it references.
    TURBO_EXPORT
    turbo::Result<Expression> replace_fields_with_known_values(const KnownFieldValues &known_values,
                                                           Expression);

    /// Simplify an expression by replacing subexpressions based on a guarantee:
    /// a boolean expression which is guaranteed to evaluate to `true`. For example, this is
    /// used to remove redundant function calls from a filter expression or to replace a
    /// reference to a constant-value field with a literal.
    TURBO_EXPORT
    turbo::Result<Expression> simplify_with_guarantee(Expression,
                                                    const Expression &guaranteed_true_predicate);

    /// Replace all named field refs (e.g. "x" or "x.y") with field paths (e.g. [0] or [1,3])
    ///
    /// This isn't usually needed and does not offer any simplification by itself.  However,
    /// it can be useful to normalize an expression to paths to make it simpler to work with.
    TURBO_EXPORT turbo::Result<Expression> remove_named_refs(Expression expression);

    /// @}

    // Execution

    /// create an ExecBatch suitable for passing to execute_scalar_expression() from a
    /// RecordBatch which may have missing or incorrectly ordered columns.
    /// Missing fields will be replaced with null scalars.
    TURBO_EXPORT turbo::Result<ExecBatch> make_exec_batch(const Schema &full_schema,
                                                        const Datum &partial,
                                                        Expression guarantee = literal(true));

    /// execute a scalar expression against the provided state and input ExecBatch. This
    /// expression must be bound.
    TURBO_EXPORT
    turbo::Result<Datum> execute_scalar_expression(const Expression &, const ExecBatch &input,
                                                 ExecContext * = nullptr);

    /// Convenience function for invoking against a RecordBatch
    TURBO_EXPORT
    turbo::Result<Datum> execute_scalar_expression(const Expression &, const Schema &full_schema,
                                                 const Datum &partial_input, ExecContext * = nullptr);

    // Serialization

    TURBO_EXPORT
    turbo::Result<std::shared_ptr<Buffer>> serialize(const Expression &);

    TURBO_EXPORT
    turbo::Result<Expression> deserialize(std::shared_ptr<Buffer>);

    /// \defgroup expression-convenience Helpers for convenient expression creation
    ///
    /// @{

    TURBO_EXPORT Expression project(std::vector<Expression> values,
                                    std::vector<std::string> names);

    TURBO_EXPORT Expression equal(Expression lhs, Expression rhs);

    TURBO_EXPORT Expression not_equal(Expression lhs, Expression rhs);

    TURBO_EXPORT Expression less(Expression lhs, Expression rhs);

    TURBO_EXPORT Expression less_equal(Expression lhs, Expression rhs);

    TURBO_EXPORT Expression greater(Expression lhs, Expression rhs);

    TURBO_EXPORT Expression greater_equal(Expression lhs, Expression rhs);

    TURBO_EXPORT Expression is_null(Expression lhs, bool nan_is_null = false);

    TURBO_EXPORT Expression is_valid(Expression lhs);

    TURBO_EXPORT Expression and_(Expression lhs, Expression rhs);

    TURBO_EXPORT Expression and_(const std::vector<Expression> &);

    TURBO_EXPORT Expression or_(Expression lhs, Expression rhs);

    TURBO_EXPORT Expression or_(const std::vector<Expression> &);

    TURBO_EXPORT Expression not_(Expression operand);

    /// @}

}  // namespace nebula::compute
